from model_info import model_urls
from torchvision.models.utils import load_state_dict_from_url
import torch.utils.model_zoo as model_zoo
class GetDatasize:
    def __init__(self):
        self.LAYER_NAME = {"inception_v3":['input', 'Conv2d_1a_3x3', 'Conv2d_2a_3x3', 'Conv2d_2b_3x3', 'Conv2d_3b_1x1', 'Conv2d_4a_3x3',
              'Mixed_5b', 'Mixed_5c', 'Mixed_5d', 'Mixed_6a', 'Mixed_6b', 'Mixed_6c', 'Mixed_6d', 'Mixed_6e',
              'Mixed_7a', 'Mixed_7b', 'Mixed_7c','prediction'],
              "resnet_101": ['input','conv1','layer_1_0', 'layer_1_1', 'layer_1_2', 'layer_2_0', 'layer_2_1', 'layer_2_2',
                            'layer_2_3', 'layer_3_0', 'layer_3_1', 'layer_3_2', 'layer_3_3', 'layer_3_4', 'layer_3_5',
                            'layer_3_6','layer_3_7','prediction'],

              "densenet121":['input','orderdict','block1','trans1','block2','trans2','block3','trans3','prediction'],
              "vgg11":['input','f0','f1','f2','f3','f4','prediction'],
              'mobilenet_v3':['input','conv1','b0', 'b1', 'b2', 'b3', 'b4', 'b5', 'b6','b7', 'b8', 'b9', 'b10', 'b11', 'b12', \
                              'b13', 'b14','conv2','prediction'],
              'convtasnet': ['input', 'group1', 'conv1', 's_0_l_0',  's_0_l_4', 's_1_l_0', 's_1_l_4',  's_2_l_0',
                             's_2_l_4', 'conv2', 'prediction']
              }
    def save_layer_weight_size(self,model_name):
        # 1. load the weight size
        if model_name in ['inception_v3','resnet_101','vgg11']:
            weight_dict = load_state_dict_from_url(model_urls[model_name])
        elif model_name == 'densenet121':
            state_dict = model_zoo.load_url(model_urls['densenet121'])
        elif model_name == 'mobilenet_v3':
            tmp = 'deeplabv3_mobilenet_v3_large_coco'
            state_dict = load_state_dict_from_url(model_urls[tmp])
        # 2. get weight by layer_name
